import os
from datetime import datetime

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from tensorflow.keras import datasets, layers,optimizers, Sequential, metrics
import io

def preposs(x,y):
    '''数据预处理'''
    x = tf.cast(x, dtype=tf.float32)/255.
    y = tf.cast(y, dtype=tf.int32)

    return x,y

# def plot_to_image(figure):
#     #save the plot to a PNG in memory
#     buf = io.BytesIO()
#     plt.savefig(buf, format='png')
#
#     plt.close(figure)
#     buf.seek(0)
#     #Convert PNG buffer to TF image
#     image = tf.image.decode_png(buf.getvalue(),channels=4)
#     #Add the batch dimension
#     image = tf.expand_dims(image, 0)
#     return image
#
# def image_grid(images):
#     '''Return a 5*5 grid of the MNIST imaages as a matplotlib figure'''
#     figure = plt.figure(figsize=(10,10))
#     for i in range(25):
#         #Start next subplot
#         plt.subplot(5, 5, i+1, title='name')
#         plt.xticks([])
#         plt.yticks([])
#         plt.grid(False)
#         plt.imshow(images[i], cmap=plt.cm.binary)

# 数据加载
(x,y),(x_test,y_test) = datasets.fashion_mnist.load_data()
print(x.shape,y.shape)

batchsz = 128 # 设置batch大小

db = tf.data.Dataset.from_tensor_slices((x,y))
db = db.map(preposs).shuffle(1000).batch(batchsz)#数据预处理、数据打散，分batch

db_test = tf.data.Dataset.from_tensor_slices((x_test,y_test))
db_test = db_test.map(preposs).shuffle(1000).batch(batchsz)#数据预处理、数据打散，分batch

db_iter = iter(db)
sample = next(db_iter)
print('batch:',sample[0].shape,sample[1].shape)

# 创建一个网络
model = Sequential([
    layers.Dense(256, activation=tf.nn.relu),
    layers.Dense(128, activation=tf.nn.relu),
    layers.Dense(64, activation=tf.nn.relu),
    layers.Dense(32, activation=tf.nn.relu),
    layers.Dense(10)
])
model.build(input_shape=[None,28*28])
model.summary()# 打印网络结构
#构建一个优化器 w = w - lr * grad
optimizer = optimizers.Adam(lr=1e-3)

acc_meter = metrics.Accuracy()
loss_meter = metrics.Mean()

# current_time = datetime.now().strftime('%Y%m%d-%H%M%S')
# log_dir = 'D:\\PycharmProjects\\DeepLearn\\图片识别\\logs'+ current_time
# summary_writer = tf.summary.create_file_writer(log_dir)
#
# # get x from (x,y)
# sample_img = next(iter(db))[0]
# sample_img =sample_img[0]
# sample_img = tf.reshape(sample_img,[1,28,28,1])
# with summary_writer.as_default():
#     tf.summary.image('Training sample:',sample_img,step=0)

def main():

    for epoch in range(30):

        for step,(x,y) in enumerate(db):

            # x: [b,28,28] => [b,784]
            # y: [b]
            x = tf.reshape(x,[-1,28*28])

            with tf.GradientTape() as tape:
                # logits:[b,10]
                logits = model(x)
                y_onehot = tf.one_hot(y,depth=10) # 将y转化为onehot
                loss_mse = tf.reduce_mean(tf.losses.MSE(y_onehot,logits))# 计算均方误差
                loss_ce = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True) # 计算交叉熵误差
                loss_ce = tf.reduce_mean(loss_ce)
                loss_meter.update_state(loss_ce)#向meter中添加loss

            grads = tape.gradient(loss_ce, model.trainable_variables)#对网络的所有可训练的参数求导
            optimizer.apply_gradients(zip(grads,model.trainable_variables))#对所有可优化参数进行更新

            if step % 100 == 0 :
                print(epoch,step,'loss:',float(loss_ce),float(loss_mse),'lossmeter:', loss_meter.result().numpy())
                loss_meter.reset_states()

        total_correct = 0
        total_num = 0
        # test
        for x,y in db_test:
            # x: [b,28,28] => [b,784]
            # y: [b]
            x = tf.reshape(x, [-1,28*28])
            #[b,10]
            logits = model(x)
            prob = tf.nn.softmax(logits, axis=1)
            indices = tf.argmax(prob, axis=1) # [b]
            indices = tf.cast(indices, dtype=tf.int32)

            result = tf.cast(tf.equal(indices, y), dtype=tf.int32) # [b]
            correct = tf.reduce_sum(result)

            total_correct += int(correct)
            total_num += x.shape[0]

        acc = total_correct/total_num

        print('epoch: acc = ', acc)

    model.save('savedmodel/model.h5')
    print('saving model.')
    pass



if __name__ == '__main__':
    main()